package com.xzq.erp.ai.memory;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.xzq.erp.ai.vo.MessageVO;
import com.xzq.erp.constants.RedisConstants;
import com.xzq.erp.ai.dto.Msg;
import lombok.Builder;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.checkerframework.checker.units.qual.A;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.ChatMemoryRepository;
import org.springframework.ai.chat.memory.InMemoryChatMemoryRepository;
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.stereotype.Component;
import org.springframework.util.Assert;

import java.util.Collections;
import java.util.List;
import java.util.Objects;

/**
 * RedisChatMemory 基于Redis实现的会话记忆存储
 */
@Slf4j
public class RedisChatMemory implements ChatMemory {

    @Autowired
    private RedisTemplate<String,String>  redisTemplate;

    @Autowired
    private ObjectMapper objectMapper;

    private static final int DEFAULT_MAX_MESSAGES = 30;

    private final ChatMemoryRepository chatMemoryRepository;

    private final int maxMessages;

    private RedisChatMemory(ChatMemoryRepository chatMemoryRepository, int maxMessages) {
        Assert.notNull(chatMemoryRepository, "chatMemoryRepository cannot be null");
        Assert.isTrue(maxMessages > 0, "maxMessages must be greater than 0");
        this.chatMemoryRepository = chatMemoryRepository;
        this.maxMessages = maxMessages;
    }

    @Override
    public void add(String conversationId, Message message) {
        log.info("RedisChatMemory add message:{}",message);
        if(Objects.isNull(message)){
            return;
        }
        //  处理会话记忆
        this.process( 1, RedisConstants.buildChatMemoryKey(conversationId));
        chatMemoryRepository.saveAll(conversationId, Collections.singletonList(message));
    }

    @Override
    public void add(String conversationId, List<Message> messages) {
         log.info("RedisChatMemory add messages:{}",messages);
         if( Objects.isNull(messages) ||  messages.isEmpty()){
             return;
         }
        //  处理会话记忆
        this.process(messages.size(), RedisConstants.buildChatMemoryKey(conversationId));
         chatMemoryRepository.saveAll(conversationId, messages);
    }

    @Override
    public List<Message> get(String conversationId) {
        log.info("RedisChatMemory get messages:{}",conversationId);
        return chatMemoryRepository.findByConversationId(conversationId);
    }

    @Override
    public void clear(String conversationId) {
        log.info("RedisChatMemory clear messages:{}",conversationId);
        chatMemoryRepository.deleteByConversationId(conversationId);
    }

    /**
     *  处理会话记忆
     * @param size 添加的会话记忆数量
     * @param key  会话记忆key
     */
    public void process(int size, String key) {
        Long chatMemorySize = redisTemplate.opsForList().size(key);
        if (Objects.nonNull(chatMemorySize)) {
            int total = (int) (chatMemorySize + size);
            if (total > maxMessages) {
                int removeCount = total - maxMessages;
                for (int i = 0; i < removeCount; i++) {
                    redisTemplate.opsForList().rightPop(key);
                }
            }
        }
    }

    public static Builder builder() {
        return new RedisChatMemory.Builder();
    }

    public static final class Builder {
        private ChatMemoryRepository chatMemoryRepository;
        private int maxMessages = DEFAULT_MAX_MESSAGES;

        private Builder() {
        }

        public Builder chatMemoryRepository(ChatMemoryRepository chatMemoryRepository) {
            this.chatMemoryRepository = chatMemoryRepository;
            return this;
        }

        public Builder maxMessages(int maxMessages) {
            this.maxMessages = maxMessages;
            return this;
        }

        public RedisChatMemory build() {
            if (this.chatMemoryRepository == null) {
                this.chatMemoryRepository = new RedisChatMemoryRepository();
            }
            log.info( "RedisChatMemory build:{}", this.maxMessages);
            return new RedisChatMemory(this.chatMemoryRepository, this.maxMessages);
        }
    }
}
